package com.tianji.aigc.service.impl;

import com.tianji.aigc.dto.ChatDTO;
import com.tianji.aigc.entity.ChatSession;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.memory.RedisChatMemory;
import com.tianji.aigc.service.ChatService;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.exceptions.BizIllegalException;
import com.tianji.common.utils.UserContext;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static org.springframework.web.servlet.function.RequestPredicates.param;

/**
 * 对话业务实现类
 *
 * @Author mr.wu
 * @Date 2025-8-17 15:28
 */
@Service
public class ChatServiceImpl implements ChatService {

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private IChatSessionService chatSessionService;

    //存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全，如果分布式环境可考虑用redis
    private static Map<String, Boolean> SESSION_MAP = new ConcurrentHashMap<String, Boolean>();

    /**
     * 聊天对话
     *
     * @param dto
     * @return
     */
    @Override
    public Flux<ChatEventVO> chat(ChatDTO dto) {
        String sessionId = dto.getSessionId();
        //判断是会话ID是否存在
        Long count = chatSessionService.lambdaQuery().eq(ChatSession::getSessionId, sessionId).count();
        if (count == 0) {
            throw new BizIllegalException("请先创建会话");
        }

        //拼接对话id：用户id + 会话id
        String conversationId = UserContext.getUser()+ ":" + sessionId;

        StringBuffer content=new StringBuffer();//线程安全的字符串拼接类，用于拼接流式输出的每个字符串

        //调用大模型进行对话
        return chatClient.prompt()
                .user(dto.getQuestion())
                .advisors(x-> param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId) )//存储对话id
                .stream()
                .content()
                .doFirst(() -> SESSION_MAP.put(sessionId, true))//标记开始输出
                .doOnError(x -> SESSION_MAP.remove(sessionId))//输出异常，清楚标记
                .doOnComplete(() -> SESSION_MAP.remove(sessionId))//标记输出完成，清楚标记
                .doOnCancel(()->{this.saveToRedis(conversationId,content.toString());})//发现取消输出时（输出停止时）将之前收集到的全部消息保存到Redis
                .takeWhile(x -> SESSION_MAP.getOrDefault(sessionId, false))//判断是否继续输出，false停止输出，true继续输出
                .map(x -> {
                    content.append(x);
                    return ChatEventVO.builder().eventData(x).eventType(ChatEventTypeEnum.DATA.getValue()).build();
                }) //大模型流式输出返回的业务数据及事件类型
                .concatWith(Flux.just(ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build())); //输出结束指定的事件类型为1002，必须为这个
    }

    @Autowired
    private RedisChatMemory redisChatMemory;

    private void saveToRedis(String conversationId, String content) {
        redisChatMemory.add( conversationId, new AssistantMessage( content));
    }

    /**
     * 停止输出
     *
     * @param sessionId
     */
    @Override
    public void stop(String sessionId) {
        SESSION_MAP.put(sessionId, false);//停止输出
    }
}